package com.gitee.wsl.security.encrypt.sample

import com.gitee.wsl.ext.base.Compare
import com.gitee.wsl.ext.map.sortByKey
import com.gitee.wsl.security.encrypt.Encrypt
import com.gitee.wsl.security.encrypt.StringToStringEncrypt

class Adfgx(val table: String,
            val keyword: String,
            val encodeMap: String = ADFGX_ENCODE_MAP,
            val replacePair: Pair<String, String> = "J" to "I") :StringToStringEncrypt{

    override fun encrypt(src: String): String =
        src.adfgx(table,keyword,encodeMap,replacePair)

    override fun decrypt(src: String): String =
        src.adfgxDecrypt(table,keyword,encodeMap)
}

fun Encrypt.Companion.adfgx(table: String,
                            keyword: String,
                            encodeMap: String = ADFGX_ENCODE_MAP,
                            replacePair: Pair<String, String> = "J" to "I") = Adfgx(table,keyword,encodeMap,replacePair)


class Adfgvx(val table: String,
            val keyword: String
) :StringToStringEncrypt{

    override fun encrypt(src: String): String =
        src.adfgvx(table,keyword)

    override fun decrypt(src: String): String =
        src.adfgvxDecrypt(table,keyword)
}

fun Encrypt.Companion.adfgvx(table: String,
                            keyword: String,
                            ) = Adfgvx(table,keyword)


const val ADFGX_ENCODE_MAP = "ADFGX"

internal fun String.adfgx(
    table: String,
    keyword: String,
    encodeMap: String = ADFGX_ENCODE_MAP,
    replacePair: Pair<String, String> = "J" to "I",
): String {
    val key = keyword.distinct()
    val polybius = polybius(table, encodeMap, replacePair)
    val keyM =
        key.fold(mutableMapOf<Char, MutableList<String>>()) { acc, c ->
            acc.apply { acc[c] = mutableListOf() }
        }
    return polybius
        .chunked(key.length)
        .fold(keyM) { acc, s ->
            acc.apply { for (i in s.indices) acc[key[i]]!!.add(s[i].toString()) }
        }.toSortedMap()
        .values
        .joinToString("") { it.joinToString("") }
}

internal fun String.adfgxDecrypt(
    table: String,
    keyword: String,
    encodeMap: String = ADFGX_ENCODE_MAP,
): String {
    val key = keyword.distinct()
    val sortedKey = key.sorted()
    val count = length % key.length
    val len = length / key.length
    val keyM2 =
        key.foldIndexed(mutableMapOf<Char, Pair<MutableList<Char>, Int>>()) { index, acc, c ->
            acc.apply {
                acc[c] = mutableListOf<Char>() to (len + (if (index < count) 1 else 0))
            }
        }.toSortedMap()
    val alphaList = this.toList().sliceList(keyM2.values.map { it.second })
    val listMap = alphaList.associateBy { sortedKey[alphaList.indexOf(it)] }
    // reverse sort
    return chunked(key.length)
        .foldIndexed(CharArray(length)) { i, acc, s ->
            acc.apply {
                s.forEachIndexed { i2, _ -> acc[i * key.length + i2] = listMap[key[i2]]!![i] }
            }
        }
        .joinToString("")
        .polybiusDecrypt(table, encodeMap)
}


const val ADFGVX_ENCODE_MAP = "ADFGVX"

internal fun String.adfgvx(table: String, keyword: String) =
    adfgx(table, keyword, ADFGVX_ENCODE_MAP, "I" to "I")

internal fun String.adfgvxDecrypt(table: String, keyword: String) =
    adfgxDecrypt(table, keyword, ADFGVX_ENCODE_MAP)


internal fun String.distinct() = asIterable().distinct().joinToString("")

internal fun String.sorted() = asIterable().sorted().joinToString("")

internal fun < V> Map<Char, V>.toSortedMap() = sortByKey(Compare.CHAR)


internal fun <T> List<T>.sliceList(split: List<Int>): MutableList<List<T>> {
    val ranges =
        split.foldIndexed(mutableListOf<IntRange>()) { index, acc, _ ->
            acc.apply {
                add(
                    split.take(index).sum() until
                            (split.take(index + 1).sum().takeIf { it <= this@sliceList.size }
                                ?: this@sliceList.size)
                )
            }
        }
    return split.indices.fold(mutableListOf()) { acc, i ->
        acc.apply { acc.add(this@sliceList.slice(ranges[i])) }
    }
}